# !usr/bin/env python
# -*- coding:utf-8 _*-

"""
@File: A2l_Updater.py

@Description:
    This file is...

@Version information:
----------------------------------------------------------------------------------------------------------------------
    Version     Date        Author  Description
----------------------------------------------------------------------------------------------------------------------
    00.00.01    2024/10/22    kook  Initial version.
"""
__VERSION__ = "00.00.01"

import re
import sys
import time
import json

from pya2l.parser import A2lParser as Parser

import ReadElf
from Logger import LOG


A2L_FILE = r"Test.a2l"
ELF_FILE = r'Test.elf'


# 数据类型映射，C语言到A2L语言的映射。updating...
mapping_c_to_a2l = {
    "char": "UBYTE",
    "unsigned char": "UBYTE",
    "signed char": "SBYTE",
    "unsigned short": "UWORD",
    "unsigned short int": "UWORD",
    "signed short": "SWORD",
    "signed short int": "SWORD",
    "unsigned long": "ULONG",
    "unsigned long int": "ULONG",
    "signed long": "SLONG",
    "signed long int": "SLONG",
    "float": "FLOAT32_IEEE",
    "double": "FLOAT64_IEEE",
}


class a2lProcess:
    def __init__(self, a2l_file, elf_file=None):
        self.a2l_file = a2l_file
        self.elf_file = elf_file
        self.parser = Parser()
        self.r_json_py = None
        self.record_dict = {}

        self._elf_vars = None
        self._elf_has_debug_info = False

        # 读取a2l文件后，再读取elf文件，更新a2l变量的地址和类型。
        self.read_a2l()
        self.update_variables()
        self.write_a2l()

    def read_a2l(self):
        # 读取a2l文件，二进制。
        with open(self.a2l_file, "rb") as f:
            byte = f.read()
            f.close()
        # a2l转换成内部的tree。
        r_ast = self.parser.tree_from_a2l(byte)
        # tree转换成json，方便修改
        r_json_bytes = self.parser.json_from_tree(r_ast)
        # load json，转换成字典，方便修改。
        self.r_json_py = json.loads(r_json_bytes.decode())

    def write_a2l(self, a2l_file="Gen.a2l"):
        # json转换成tree。
        w_json = json.dumps(self.r_json_py, indent=4)
        w_ast = self.parser.tree_from_json(w_json.encode())
        # tree转换成a2l。
        w_a2l = self.parser.a2l_from_tree(w_ast, indent=4).decode()  # SYMBOL_LINK没了
        w_a2l = self.add_SYMBOL_LINK(w_a2l)
        print(w_a2l, file=open(a2l_file, "w"))
        print(f"Save file: {a2l_file}")

    @staticmethod
    def add_SYMBOL_LINK(string):
        """
        变量添加SYMBOL_LINK。
        tree转成的a2l丢失了SYMBOL_LINK，这个是Vector工具更新变量地址用到的。
        """
        regex = re.compile(r"(/begin (MEASUREMENT|CHARACTERISTIC) ([^ ]+) .*)")
        return regex.sub(r"\1" + " SYMBOL_LINK " + r'"\3" 0', string)

    def read_record(self):
        """
        a2l中的RECORD_LAYOUT，数据类型。
        """
        project = self.r_json_py.get("PROJECT")
        modules = project.get("MODULE")
        for module in modules:  # MODULE长度是1
            records = module.get("RECORD_LAYOUT")
            for record in records:
                name = record["Name"]["Value"]
                fnc = record.get("FNC_VALUES") or record.get("AXIS_PTS_X")
                data_type = fnc["DataType"]["Value"]
                if not self.record_dict.get(data_type):
                    self.record_dict[data_type] = {data_type}
                self.record_dict[data_type].add(name)
        # print(f"record_dict: {self.record_dict}")

    def update_measurements(self):
        """
        更新观测量的地址和类型。
        """
        project = self.r_json_py.get("PROJECT")
        modules = project.get("MODULE")
        for module in modules:  # MODULE长度是1
            measurements = module.get("MEASUREMENT")
            for variable in measurements:
                # a2l中变量的属性
                name = variable["Name"]["Value"]
                data_type = variable["DataType"]["Value"]
                address = int(variable["ECU_ADDRESS"]["Address"].get("Value", "0"))
                # elf中变量的属性
                ret = self.get_elf_var_info(name)
                if ret is None:
                    continue
                new_addr, elf_dt = ret
                if new_addr == address:
                    LOG.info(f"{name} address is the same as before.")
                else:
                    variable["ECU_ADDRESS"]["Address"]["Value"] = new_addr
                    LOG.info(f"{name} address is updated from {hex(address)} to {hex(new_addr)}.")
                # elf变量数据类型转换成a2l数据类型。
                if self._elf_has_debug_info:
                    new_dt = mapping_c_to_a2l.get(elf_dt)
                    if new_dt is None:
                        LOG.error(f"Unknown C language dataType {new_dt}, please add in mapping_c_to_a2l. {name}")
                        continue
                    # 对比原始a2l变量类型和计算出的数据类型是否相同，不同则更新数据类型。
                    a2l_dt_set = self.record_dict.get(new_dt)
                    if a2l_dt_set is not None:
                        if data_type not in a2l_dt_set:
                            variable["DataType"]["Value"] = iter(a2l_dt_set).__next__()
                            LOG.info(f"{name} data_type is updated from {data_type} to {iter(a2l_dt_set).__next__()}.")
                    else:
                        LOG.error(f"{new_dt} is not found in record_dict, update a2l RECORD_LAYOUT.{self.record_dict}")

    def update_characteristics(self):
        """
        更新标定量的地址和类型。
        """
        project = self.r_json_py.get("PROJECT")
        modules = project.get("MODULE")
        for module in modules:  # MODULE长度是1
            characteristics = module.get("CHARACTERISTIC")
            for variable in characteristics:
                name = variable["Name"]["Value"]
                data_type = variable["Deposit"]["Value"]
                address = int(variable["Address"].get("Value", "0"))

                ret = self.get_elf_var_info(name)
                if ret is None:
                    continue
                new_addr, elf_dt = ret
                if new_addr == address:
                    LOG.info(f"{name} address is the same as before.")
                else:
                    variable["Address"]["Value"] = new_addr
                    LOG.info(f"{name} address is updated from {hex(address)} to {hex(new_addr)}.")

                # elf变量数据类型转换成a2l数据类型。
                if self._elf_has_debug_info:
                    new_dt = mapping_c_to_a2l.get(elf_dt)
                    if new_dt is None:
                        LOG.error(f"Unknown C language dataType {new_dt}, please add in mapping_c_to_a2l. {name}")
                        continue
                    a2l_dt_set = self.record_dict.get(new_dt)
                    if a2l_dt_set is not None:
                        if data_type not in a2l_dt_set:
                            variable["Deposit"]["Value"] = iter(a2l_dt_set).__next__()
                            LOG.info(f"{name} data_type is updated from {data_type} to {iter(a2l_dt_set).__next__()}.")
                    else:
                        LOG.error(f"{new_dt} is not found in record_dict, update a2l RECORD_LAYOUT.{self.record_dict}")

    def update_variables(self):
        """
        更新变量。
        """
        self.read_record()
        self.update_measurements()
        self.update_characteristics()

    @property
    def elf_vars(self):
        if self._elf_vars is None:
            try:
                elf_parser = ReadElf.Parser(self.elf_file)
                _ = elf_parser.read_variables()  # 加个返回值，避免打印顺序错乱。
                self._elf_vars = elf_parser.vars
                self._elf_has_debug_info = elf_parser.has_debug_info()
                if not self._elf_has_debug_info:
                    LOG.critical(f"No debug info in {self.elf_file}, "
                                 f"so can't get variable data type and struct/array address.")
            except:
                LOG.error(f"Unknown elf file. {self.elf_file}")
                exit(-1)
        return self._elf_vars

    def get_elf_var_info(self, var):
        """
        获取elf中变量的地址和数据类型。
        :param var: a2l中变量的名称。包含结构体和数组类型的变量。
        :return: (int, str) 变量地址，变量C语言数据类型。
        """
        var_splits = var.split(".")
        if not self._elf_has_debug_info and len(var_splits) > 1:
            return None  # 结构体和数组变量，需要debug信息才能获取地址和数据类型。

        base_var = var_splits[0]  # 可能是结构体或者数组，取变量的第一级。
        info = self.elf_vars.get(base_var)  # type: ReadElf.Item

        if info is None:
            LOG.error(f"{var} is not found in elf_vars.")
            return None
        addr = info.addr
        elf_dt = info.data_type  # type: str | ReadElf.Item

        # 结构体或数组变量，继续搜索下一级。
        if len(var_splits) > 1:
            for x in var_splits[1:]:
                regex = re.compile(r"_(\d+)_")
                mo = regex.search(x)
                if mo:  # array
                    index = int(mo.group(1))
                    if index < elf_dt.array_length:
                        addr += index * elf_dt.data_type_size
                    else:
                        LOG.error(f"Index {index} out of range for {var}.")
                        return None
                else:
                    for child in elf_dt.children:  # type: ReadElf.Item
                        if child.name == x:
                            elf_dt = child
                            addr += child.member_offset
                            break
                    else:
                        LOG.error(f"{x} is not found in {var}.")
                        return None
        # 不是str，则是Item，即结构体类型。结构体类型直接返回unsigned char，a2l中对应UBYTE。
        if isinstance(elf_dt, ReadElf.Item):
            elf_dt = elf_dt.data_type
        LOG.debug(f"{var} addr: {hex(addr)}, data_type: {elf_dt}")

        return addr, elf_dt


def main():
    t = time.time()
    global A2L_FILE
    global ELF_FILE
    for x in sys.argv:
        if x.lower().endswith('.elf'):
            ELF_FILE = x
        if x.lower().endswith('.a2l'):
            A2L_FILE = x
    LOG.critical(f"{A2L_FILE=} {ELF_FILE=}")
    a2lProcess(A2L_FILE, ELF_FILE)
    print(f"Executed time: {time.time() - t}s")


if __name__ == '__main__':
    main()
